/**
 * @file main.cpp
 *
 * Copyright (C) 2024. Huawei Technologies Co., Ltd. All rights reserved.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
 */
#include "add_custom_tiling.h"
#include "data_utils.h"
#ifndef ASCENDC_CPU_DEBUG
#include "acl/acl.h"
#include "aclrtlaunch_add_custom.h"
#else
#include "tikicpulib.h"
extern "C" __global__ __aicore__ void add_custom(GM_ADDR x, GM_ADDR y, GM_ADDR z, AddCustomTilingData tiling);
#endif

int32_t main(int32_t argc, char *argv[])
{
    uint32_t blockDim = 48;
    uint32_t mBlock = 128;
    uint32_t n = 2048;
    uint32_t m = blockDim * mBlock;
    size_t len = m * n;
    size_t tilingSize = 4 * sizeof(uint32_t);
    size_t inputAByteSize = len * sizeof(uint16_t);
    size_t inputXByteSize = n * sizeof(uint16_t);
    size_t outputYByteSize = m * sizeof(uint16_t);

    CHECK_ACL(aclInit(nullptr));
    int32_t deviceId = 0;
    CHECK_ACL(aclrtSetDevice(deviceId));
    aclrtStream stream = nullptr;
    CHECK_ACL(aclrtCreateStream(&stream));

    AddCustomTilingData *tiling;
    uint8_t *AHost, *XHost, *YHost;
    uint8_t *ADevice, *XDevice, *YDevice;

    CHECK_ACL(aclrtMallocHost((void **)(&tiling), tilingSize));
    ReadFile("./input/input_tiling.bin", tilingSize, tiling, tilingSize);

    CHECK_ACL(aclrtMallocHost((void **)(&AHost), inputAByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&XHost), inputXByteSize));
    CHECK_ACL(aclrtMallocHost((void **)(&YHost), outputYByteSize));
    CHECK_ACL(aclrtMalloc((void **)&ADevice, inputAByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&XDevice, inputXByteSize, ACL_MEM_MALLOC_HUGE_FIRST));
    CHECK_ACL(aclrtMalloc((void **)&YDevice, outputYByteSize, ACL_MEM_MALLOC_HUGE_FIRST));

    ReadFile("./input/matrix_A.bin", inputAByteSize, AHost, inputAByteSize);
    ReadFile("./input/vector_X.bin", inputXByteSize, XHost, inputXByteSize);

    CHECK_ACL(aclrtMemcpy(ADevice, inputAByteSize, AHost, inputAByteSize, ACL_MEMCPY_HOST_TO_DEVICE));
    CHECK_ACL(aclrtMemcpy(XDevice, inputXByteSize, XHost, inputXByteSize, ACL_MEMCPY_HOST_TO_DEVICE));

    for (int i = 0; i < 100; ++i)
    {
        ACLRT_LAUNCH_KERNEL(add_custom)(blockDim, stream, ADevice, XDevice, YDevice, tiling);
    }

    int num_repeat = 100;

    aclrtEvent start, stop;
    float temp_time = 0;
    float time = 0;
    CHECK_ACL(aclrtCreateEvent(&start));
    CHECK_ACL(aclrtCreateEvent(&stop));
    for (int i = 0; i < num_repeat; ++i)
    {
        CHECK_ACL(aclrtSynchronizeStream(stream));
        CHECK_ACL(aclrtRecordEvent(start, stream));

        ACLRT_LAUNCH_KERNEL(add_custom)(blockDim, stream, ADevice, XDevice, YDevice, tiling);

        CHECK_ACL(aclrtSynchronizeStream(stream));
        CHECK_ACL(aclrtRecordEvent(stop, stream));
        CHECK_ACL(aclrtSynchronizeEvent(stop));
        CHECK_ACL(aclrtEventElapsedTime(&temp_time, start, stop));
        time += temp_time;
    }

    printf("m: %u, n: %u, %f GFLOPS, %f ms\n", m, n,
           (float)2 * m * n / (time / num_repeat * 1e-3) / 1e9, (time / num_repeat));

    CHECK_ACL(aclrtMemcpy(YHost, outputYByteSize, YDevice, outputYByteSize, ACL_MEMCPY_DEVICE_TO_HOST));
    WriteFile("./output/output_Y.bin", YHost, outputYByteSize);

    CHECK_ACL(aclrtFree(ADevice));
    CHECK_ACL(aclrtFree(XDevice));
    CHECK_ACL(aclrtFree(YDevice));
    CHECK_ACL(aclrtFreeHost(AHost));
    CHECK_ACL(aclrtFreeHost(XHost));
    CHECK_ACL(aclrtFreeHost(YHost));
    CHECK_ACL(aclrtFreeHost(tiling));

    CHECK_ACL(aclrtDestroyStream(stream));
    CHECK_ACL(aclrtResetDevice(deviceId));
    CHECK_ACL(aclFinalize());

    return 0;
}
